import xml.dom.minidom as xmldom
import os
import collections


file_loc="./VOC2007/train_val/VOCdevkit/VOC2007/Annotations"
file_list = os.listdir(file_loc) 
file_list =[i for i in file_list if i.endswith('.xml')]
#print(file_list)


ObjectClass = collections.namedtuple('ObjectClass',['index', 'name', 'xmin', 'ymin', 'xmax', 'ymax'])
PicClass = collections.namedtuple('PicClass',['index', 'width', 'height', 'object_list'])


name_list=['diningtable', 'motorbike', 'bicycle', 'pottedplant', 'aeroplane', 'train', 'car', 'dog', 'bottle', 'chair', 'cat', 'cow', 'sheep', 'boat', 'sofa', 'person', 'tvmonitor', 'bird', 'horse', 'bus']


def all_files():
    file_obj_list=[]

    for _file in file_list:  #['000648.xml']:
         _file=_file[:-4]
         #all xml files, list set of PicObject
         file_obj_list.append( single_file(_file) )

    return file_obj_list


def single_file(file_index_str):
    # file_index_str "000648"
    xmlfilepath=file_loc+"/"+file_index_str+".xml"
    domobj = xmldom.parse(xmlfilepath)
    elementobj = domobj.documentElement

    #find object elements
    subElementObj=elementobj.getElementsByTagName("object")
    #print(subElementObj)

    #the  width element of object element
    width=elementobj.getElementsByTagName("width")[0].firstChild.data
    #the  height element of object element
    height=elementobj.getElementsByTagName("height")[0].firstChild.data

    #the object list set of per picture
    per_pic_obj_list=[]

    for element in subElementObj:
        #print( subElementObj[0].getElementsByTagName("name")[0] )
        #print( subElementObj[0].getElementsByTagName("name")[0].firstChild )
        name = element.getElementsByTagName("name")[0].firstChild.data 
        xmin = element.getElementsByTagName("xmin")[0].firstChild.data 
        ymin = element.getElementsByTagName("ymin")[0].firstChild.data
        xmax = element.getElementsByTagName("xmax")[0].firstChild.data
        ymax = element.getElementsByTagName("ymax")[0].firstChild.data
  
        # one object which has four elements,such as name,xmin,ymin,xmax,ymax
        obj = ObjectClass(name_list.index(name), name, xmin, ymin, xmax, ymax)

        #add object into list which belong to one picture
        per_pic_obj_list.append(obj)

    # picture class for one picture     
    picClass=PicClass(file_index_str, width, height, per_pic_obj_list)

    return picClass


if __name__=="__main__":
    #test one picture's xml which name is 000648
    print( single_file("000648") )

    print( all_files() )










